
from pydantic import Field
from core.llm.base_analyze import BaseAnalyzer
from core import prompt_dir
from typing import Dict
from loguru import logger


class RewriteQueryAnalyzer(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / '问题改写.md')

    def invoke(
            self,
            userinfo: Dict = Field(description='用户信息'),
            schema: Dict = Field(description="数据表结构描述"),
            query: str = Field(description='用户的问题描述')
    ):
        schema.pop('rows', [])
        prompt = self.template_prompt.format(
            inputs={
                "userinfo": userinfo,
                "data_info": schema,
                "query": query
            },
            remove_template_variables=True
        )
        data = self.analyze(prompt)
        name = schema.get('name', '')
        logger.info(f'表：{name}  --- 问题改写：{data}')
        return data
